import argparse
from datetime import datetime,timedelta
import os.path as osp
import time
import cv2
import torch
import torch.utils.data
from tqdm import tqdm
from datasets import build_test_loader, build_train_loader,build_train_loader_dataset,build_train_loader_mlc,build_test_loader_ssm
from defaults import get_default_cfg
from engine_mlc import evaluate_performance, train_one_epoch,train_one_epoch_prw_with_ssm
from models.seqnet_mlc import SeqNet
from utils.utils import mkdir, resume_from_ckpt, save_on_master, set_random_seed,resume_from_ckpt_model_only
import pdb
import numpy as np
from utils.utils import Logger,resume_from_ckpt_model_only_psmlc
import sys
from models.memory import Memory
from models.mmcl import MMCL
from demo import detector,detector_box_only,detector_box_only_ssm
from utils.utils import unpickle,pickle
from scipy.io import loadmat
from datasets.build import print_statistics_mlc
from pandas import DataFrame
from init_mem_for_boxes import *

def main(args):
    #------------------------------配置参数---------------------------------------------------
    cfg = get_default_cfg()
    if args.cfg_file:
        cfg.merge_from_file(args.cfg_file)
    cfg.merge_from_list(args.opts)
    if cfg.single_label:
        cfg.NEG_RATE=1.0
    else:
        cfg.NEG_RATE=float(args.neg_rate)
    cfg.score_anno=args.score_anno
    cfg.THRE_SIM=float(args.sim_thresh)
    cfg.delta=int(args.delta)
    cfg.freeze()
    device = torch.device(cfg.DEVICE)
    if cfg.SEED >= 0:
        set_random_seed(cfg.SEED)

    if not args.eval:
        sys.stdout = Logger(osp.join(cfg.OUTPUT_DIR, 'log_train.txt'))
    else:
        sys.stdout = Logger(osp.join(cfg.OUTPUT_DIR, 'log_test.txt'))
    print("==========\nArgs:{}\n==========".format(args))
    for i in str(args).split(','):
        print(i)
    print("====================")
    print(cfg.NEG_RATE)
    #------------------------------加载数据，对train_loader后续更改---------------------------------------------------
    print("Loading org data ")
    train_loader,dataset= build_train_loader_dataset(cfg)
    gallery_loader, query_loader = build_test_loader(cfg)

    #------------------------------模型加载---------------------------------------------------
    print("Creating model")
    model = SeqNet(cfg) # wu oim
    model.to(device)
    result=[]

    if args.eval:
        assert args.ckpt, "--ckpt must be specified when --eval enabled"
        resume_from_ckpt_model_only_psmlc(args.ckpt, model)
        eval_all,eval_labeled,eval_reid=evaluate_performance(
            model,
            gallery_loader,
            query_loader,
            device,
            use_gt=cfg.EVAL_USE_GT,
            use_cache=cfg.EVAL_USE_CACHE,
            use_cbgm=cfg.EVAL_USE_CBGM,
        )
        labeled_only,count_tp,count_gt,det_rate,y_true,ap=eval_all
        print("{} detection:".format("all"))
        print("  recall=TP/all_positive ={}/{}= {:.2%}".format(count_tp,count_gt,det_rate))
        print("  presicion=TP/(TP+FP)={}/{} = {:.2%}".format(count_tp,len(y_true),count_tp/(len(y_true))))
        print("  ap = {:.2%}".format(ap))
        labeled_only,count_tp,count_gt,det_rate,y_true,ap=eval_labeled
        print("{} detection:".format("labeled only"))
        print("  recall=TP/all_positive ={}/{}= {:.2%}".format(count_tp,count_gt,det_rate))
        print("  presicion=TP/(TP+FP)={}/{} = {:.2%}".format(count_tp,len(y_true),count_tp/(len(y_true))))
        print("  ap = {:.2%}".format(ap))
        aps,accs,topk=eval_reid
        print("search ranking:")
        print("  mAP = {:.2%}".format(np.mean(aps)))
        accs = np.mean(accs, axis=0)
        for i, k in enumerate(topk):
            print("  top-{:2d} = {:.2%}".format(k, accs[i]))

        reidmAP_r = np.mean(eval_reid[0]),
        reidtop1 = accs[0],
        alldetrecall = eval_all[3],
        alldetpresicion = eval_all[1] / len(eval_all[4]),
        alldetap = eval_labeled[-1],
        labdetrecall = eval_labeled[3],
        labdetpresicion = eval_labeled[1] / len(eval_labeled[4]),
        labdetap = eval_labeled[-1],

        res = [reidmAP_r, reidtop1, alldetrecall, alldetpresicion, alldetap, labdetrecall, labdetpresicion,labdetap]

        res = [round(x[0], 4) * 100 for x in res]
        result.append(res)

        arr = np.array(result)
        df = DataFrame(arr)
        df.columns = ['mAP', "top1", "all-det", "all-pre", "all-ap", "lab-det", "lab-pre", "lab-ap"]
        df.to_csv(cfg.OUTPUT_DIR + '/reid_res.csv')
        print(df)
        sys.exit()

    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(
        params,
        lr=cfg.SOLVER.BASE_LR,
        momentum=cfg.SOLVER.SGD_MOMENTUM,
        weight_decay=cfg.SOLVER.WEIGHT_DECAY,
    )

    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=cfg.SOLVER.LR_DECAY_MILESTONES, gamma=cfg.SOLVER.GAMMA
    )

    start_epoch = args.start_epoch
    if args.resume:
        assert args.ckpt, "--ckpt must be specified when --resume enabled"
        if args.init_cd==True:
            resume_from_ckpt_model_only_psmlc(args.ckpt, model, optimizer, lr_scheduler)
        else:
            start_epoch = resume_from_ckpt(args.ckpt, model, optimizer, lr_scheduler)+1
        # start_epoch = resume_from_ckpt_model_only(args.ckpt, model, optimizer, lr_scheduler) + 1

    path = osp.join(cfg.OUTPUT_DIR, "config.yaml")
    with open(path, "w") as f:
        f.write(cfg.dump())
    print(f"Full config is saved to {path}")

    tfboard = None
    if cfg.TF_BOARD:
        from torch.utils.tensorboard import SummaryWriter
        timestamp = "{0:%Y-%m-%dT%H-%M-%S}".format(datetime.now())
        tf_log_path = osp.join(cfg.OUTPUT_DIR, timestamp + "tf_log")
        mkdir(tf_log_path)
        tfboard = SummaryWriter(log_dir=tf_log_path)
        print(f"TensorBoard files are saved to {tf_log_path}")

    print("Start training")

    best_reidmAP=0
    best_epoch=0
    criterion = MMCL(cfg.delta, cfg.NEG_RATE,cfg).to(device)  # cfg.MMCL.DELTA, cfg.MMCL.R,cfg.Device

    for epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCHS):
        if epoch==0: # 初始化MEM特征 都为 0
            print('gt or detect new box for mplp...')  #先GT，后 yolo
            if not cfg.gtBoxes:# un-mlc 和transfer只在create dataset有区别 其他没有区别  已经mplp的epoch2使用single class id
                if cfg.INPUT.DATASET=='CUHK-SYSU':
                    annotations, memory=create_new_pid_ssm_yolo5l(cfg)
                else: #prw-un-mlc with yolo-boxes
                    annotations, memory=create_new_pid_prw(cfg)
            else: # use gtBoxes
                if cfg.INPUT.DATASET=='CUHK-SYSU':
                    annotations, memory=create_new_pid_ssm_gt(cfg,dataset)
                else: #prw
                    annotations, memory=create_new_pid_ssm_gt(model,cfg)
            new_loader = build_train_loader_mlc(cfg,annotations)

        train_one_epoch(cfg, model, memory,criterion,optimizer, new_loader, device, epoch, tfboard,
                        gt_multilabel_source=None)

        # torch.save(memory.lut,'../lut_epoch_'+str(epoch)+'.pth')

        if tfboard:
            lr=lr_scheduler.get_last_lr()[0]
            print('epoch: ',epoch,'lr :',lr)
            # tfboard.add_scalars('lr',{'lr-epoch':lr},epoch)
        lr_scheduler.step()

        #ckpt save
        if (epoch + 1) % cfg.CKPT_PERIOD == 0 or epoch == cfg.SOLVER.MAX_EPOCHS - 1:
            save_on_master(
                {
                    "model": model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "lr_scheduler": lr_scheduler.state_dict(),
                    "epoch": epoch,
                },
                osp.join(cfg.OUTPUT_DIR, str(epoch)+'epoch.pth'),
            )
        #eval 
        if (epoch + 1) % cfg.EVAL_PERIOD == 0 or epoch == cfg.SOLVER.MAX_EPOCHS - 1:
            eval_all,eval_labeled,eval_reid=evaluate_performance(
                model,
                gallery_loader,
                query_loader,
                device,
                use_gt=cfg.EVAL_USE_GT,
                use_cache=cfg.EVAL_USE_CACHE,
                use_cbgm=cfg.EVAL_USE_CBGM,
            )
            # eval_all=(1,2,3,4,[56,7,8],6)
            # eval_labeled=(4,5,6,7,[6.7,8],9)
            # eval_reid=([2,3,4,6],[[2,3,4],[2,3,4],[2,3,4]],[1,5,10])
            labeled_only,count_tp,count_gt,det_rate,y_true,ap=eval_all
            print("{} detection:".format("all"))
            print("  recall=TP/all_positive ={}/{}= {:.2%}".format(count_tp,count_gt,det_rate))
            print("  presicion=TP/(TP+FP)={}/{} = {:.2%}".format(count_tp,len(y_true),count_tp/(len(y_true))))
            print("  ap = {:.2%}  FP={}".format(ap,len(y_true)-count_tp))
            labeled_only,count_tp,count_gt,det_rate,y_true,ap=eval_labeled
            print("{} detection:".format("labeled only"))
            print("  recall=TP/all_positive ={}/{}= {:.2%}".format(count_tp,count_gt,det_rate))
            print("  presicion=TP/(TP+FP)={}/{} = {:.2%}".format(count_tp,len(y_true)-count_tp,count_tp/(len(y_true))))
            print("  ap = {:.2%}".format(ap))
            aps,accs,topk=eval_reid
            print("search ranking:")
            print("  mAP = {:.2%}".format(np.mean(aps)))
            accs = np.mean(accs, axis=0)
            for i, k in enumerate(topk):
                print("  top-{:2d} = {:.2%}".format(k, accs[i]))
            reidmAP=np.mean(aps)
        #ckpt save best
        if (epoch + 1) % cfg.EVAL_PERIOD == 0 or epoch == cfg.SOLVER.MAX_EPOCHS - 1:
            if reidmAP>best_reidmAP:
                best_reidmAP=reidmAP
                best_epoch=epoch
                save_on_master(
                    {
                        "model": model.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "lr_scheduler": lr_scheduler.state_dict(),
                        "epoch": best_epoch,
                    },
                    osp.join(cfg.OUTPUT_DIR, 'best.pth'),
                )
        if tfboard and (epoch + 1) % cfg.EVAL_PERIOD == 0:#  aps,accs,topk=eval_reid
        #labeled_only,count_tp,count_gt,det_rate,y_true,ap=eval_lab
            tfboard.add_scalars('eval',
            {'reid-mAP':np.mean(eval_reid[0]),
             'reid-top1':accs[0],
             'reid-top5':accs[1],
             'reid-top10':accs[2],
             'all det-recall':eval_all[3],
             'all det-presicion':eval_all[1]/len(eval_all[4]),
             'all det-ap':eval_all[-1],
             'lab det-recall': eval_labeled[3],
             'lab det-presicion': eval_labeled[1] / len(eval_labeled[4]),
             'lab det-ap': eval_labeled[-1],
             },epoch)
            reidmAP_r = np.mean(eval_reid[0]),
            reidtop1 = accs[0],
            alldetrecall = eval_all[3],
            alldetpresicion = eval_all[1] / len(eval_all[4]),
            alldetap = eval_all[-1],
            labdetrecall = eval_labeled[3],
            labdetpresicion = eval_labeled[1] / len(eval_labeled[4]),
            labdetap = eval_labeled[-1],

            res = [reidmAP_r, reidtop1, alldetrecall, alldetpresicion, alldetap, labdetrecall, labdetpresicion,
                   labdetap]
            res = [round(x[0], 4) * 100 for x in res]
            result.append(res)

            arr = np.array(result)
            df = DataFrame(arr)
            df.columns = ['mAP', "top1", "all-det", "all-pre", "all-ap", "lab-det", "lab-pre", "lab-ap"]
            df.to_csv(cfg.OUTPUT_DIR + '/reid_res.csv')
            print(df)
    
    if tfboard:
        tfboard.close()
    print('best mAP{}, best epoch: {}'.format(best_reidmAP,best_epoch))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train a person search network.")
    parser.add_argument("--cfg",dest='cfg_file',default="configs/1ssm_un_single_yolov5l.yaml", help="Path to configuration file.")
                            #           "configs/ssm_un_mlc.yaml  
    parser.add_argument("--init_cd", default=True)#action="store_true", help="resume and init_cd ckpt.")
    parser.add_argument("--eval", action="store_true", help="Evaluate the performance of a given checkpoint.")
    parser.add_argument("--resume", default=False)#action="store_true", help="Resume from the specified checkpoint.")
    parser.add_argument("--start_epoch", default=0,type=int)
    parser.add_argument("--sim_thresh", default=0.6)

    train_prw=False
    if train_prw:
        parser.add_argument("--data_name", default='prw')  # prw
        parser.add_argument("--ckpt",default=None) #'../author_model/epoch_19.pth')  # help="Path to checkpoint to resume or evaluate.")
        parser.add_argument("--neg_rate", default=0.01)  # ssm 0.001 * 60000 = 600, 但是互近邻法则会提出假正例
        parser.add_argument("--delta", default=6)
        parser.add_argument("--score_anno", default='data/extract_prw_anno_yolo0.0_28624.pkl')

    else:
        parser.add_argument("--data_name", default='ssm')  
        parser.add_argument("--ckpt",default=None) #'../author_model/epoch_17_prw.pth')# help="Path to checkpoint to resume or evaluate.")
        parser.add_argument("--neg_rate", default=0.001)  # 先设置为1 # ssm 0.001 * 60000 = 600, 但是互近邻法则去除假正例
        parser.add_argument("--delta", default=5)
        parser.add_argument("--score_anno", default='data/extra_ssm_yolo_0.0_65020.json')

    parser.add_argument("opts", nargs=argparse.REMAINDER, help="Modify config options using the command-line")
    args = parser.parse_args()
    main(args)
